Skip to content

feat(ce,flce): decouple gradients computation for no_grad mode #894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 10 commits into from
Oct 11, 2025
Merged

feat(ce,flce): decouple gradients computation for no_grad mode #894

merged 10 commits into from
Oct 11, 2025

Conversation

@Tcc0403
Copy link
Collaborator

@Tcc0403 Tcc0403 commented Oct 1, 2025 •
edited
Loading

Summary

Add a flag HAS_GRADIENTS to cross entropy kernel. No more gradients computation if there's no need.

Testing Done

Cross Entropy forward with no_grad
image

Fused Linear Cross Entropy forward with no_grad
image

  • Hardware Type:
  • run make test to ensure correctness
  • run make checkstyle to ensure code style
  • run make test-convergence to ensure convergence

@Tcc0403 Tcc0403 changed the title feat(ce,-flce): decouple gradients computation for no_grad mode feat(ce,flce): decouple gradients computation for no_grad mode Oct 1, 2025
@Tcc0403 Tcc0403 force-pushed the tcc/flce-eval-no-grad branch from 291141c to 9ab603a Compare October 1, 2025 13:40
Tcc0403 added 4 commits October 1, 2025 21:41
@Tcc0403
format
8761844 
Signed-off-by: Tcc0403 <[email protected]>
@Tcc0403
fix grad calc condition
f1f525b 
Signed-off-by: Tcc0403 <[email protected]>
@Tcc0403
fix ce test
bd78d85 
Signed-off-by: Tcc0403 <[email protected]>
@Tcc0403
check _input instead of _input_chunk
4fa3179 
Signed-off-by: Tcc0403 <[email protected]>
@Tcc0403 Tcc0403 mentioned this pull request Oct 3, 2025
Merged
3 tasks
@Tcc0403 Tcc0403 requested review from momochen and shimizust October 8, 2025 10:53
@Tcc0403 Tcc0403 mentioned this pull request Oct 8, 2025
Closed
3 tasks
shimizust
shimizust approved these changes Oct 11, 2025
View reviewed changes
Copy link
Collaborator

@shimizust shimizust left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, thanks for adding this. The forward flce still is significantly slower than hf since we're still computing grad_input applying token scaling logic? Also do you know why fp32 accum is faster?

@lancerts lancerts merged commit 5c2a04d into main Oct 11, 2025
2 checks passed
@lancerts lancerts deleted the tcc/flce-eval-no-grad branch October 11, 2025 15:44
@Tcc0403
Copy link
Collaborator Author

Tcc0403 commented Oct 11, 2025 •
edited
Loading

@shimizust The slower forward pass is kinda expected because:

  1. instead of one big matmul, we slice input rowwise and do multiple matmuls. It brings some kernel launch overhead. Plus, it's not gauranteed to ivnoke the most efficient kernels (tiling size, tail effect, etc... from kernel's prespective).
  2. multiple cross entropy kernel launch overhead and similar issue also add up.

I just found that we can remove this line in eval mode as well, cutting another matmul for each interation should be significant.

Liger-Kernel/src/liger_kernel/ops/fused_linear_cross_entropy.py

Line 175 in 5c2a04d

grad_input[start_idx:end_idx] = grad_logits_chunk @ weight

and some grad tensors allocations too

@Tcc0403 Tcc0403 mentioned this pull request Oct 11, 2025
Merged
3 tasks
shimizust pushed a commit that referenced this pull request Oct 16, 2025
@Tcc0403
Filter out redundant ops/allocations in no_grad mode (#906)
c856fba 
## Summary
<!--- This is a required section; please describe the main purpose of
this proposed code change. --->
follow-up #894 


<!---
## Details
This is an optional section; is there anything specific that reviewers
should be aware of?
--->

## Testing Done
<!--- This is a required section; please describe how this change was
tested. --->

<!-- 
Replace BLANK with your device type. For example, A100-80G-PCIe

Complete the following tasks before sending your PR, and replace `[ ]`
with
`[x]` to indicate you have done them. 
-->

- Hardware Type: <BLANK>
- [ ] run `make test` to ensure correctness
- [ ] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Signed-off-by: Tcc0403 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@shimizust shimizust shimizust approved these changes

@momochen momochen Awaiting requested review from momochen

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@Tcc0403 @shimizust @lancerts